/**
 * @file syncbntrainingupdate_custom_tiling.cpp
 *
 * Copyright (C) 2025. Huawei Technologies Co., Ltd. All rights reserved.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.
 */

#include <cassert>
#include "syncbntrainingupdate_custom_tiling.h"

// 数据类型大小映射 (float32)
constexpr uint32_t DATA_TYPE_SIZE = 4;
constexpr uint32_t BLOCK_SIZE = 32;
constexpr uint32_t BUFFER_NUM = 2;  // 双缓冲
constexpr uint32_t UB_BLOCK_NUM = 100;  // UB最大可以使用的block数量
constexpr uint32_t MAX_AVAILABLE_UB_BLOCK_NUM = UB_BLOCK_NUM / BUFFER_NUM * BUFFER_NUM;
// Tiling参数计算函数
void TilingParamsCalc(uint32_t length, uint32_t ubBlockNum,
    uint32_t& tileNum, uint32_t& tileLength, uint32_t& lastTileLength)
{
    assert(ubBlockNum != 0U);
    tileNum = length / ubBlockNum;
    if (length % ubBlockNum == 0U || tileNum == 0U) {
        if (tileNum == 0U) {
            tileNum = 1U;
        }
        if (length < ubBlockNum) {
            tileLength = length;
            lastTileLength = tileLength;
        } else {
            tileLength = ubBlockNum;
            lastTileLength = tileLength;
        }
    } else {
        tileNum++;
        tileLength = ubBlockNum;
        lastTileLength = (uint32_t)(length - (tileNum - 1) * tileLength);
    }
}


 void GenerateTilingData(uint8_t* tilingBuf, uint32_t blockDim)
{

    SyncBNTrainingUpdateCustomTilingData *tiling = reinterpret_cast<SyncBNTrainingUpdateCustomTilingData *>(tilingBuf);
    // 计算对齐参数
    uint32_t alignNum = BLOCK_SIZE / DATA_TYPE_SIZE;
    assert(alignNum != 0U);

    // 如果已经设置了totalLength，则计算分块参数
    if (tiling->totalLength > 0) {
        // 计算对齐参数
        uint32_t alignNum = BLOCK_SIZE / DATA_TYPE_SIZE;
        assert(alignNum != 0U);
        
        // 计算UB中可用的块数（考虑6个缓冲区）
        uint32_t ubBlockAligned = MAX_AVAILABLE_UB_BLOCK_NUM * alignNum / (6 * BUFFER_NUM) * (6 * BUFFER_NUM);
        if (ubBlockAligned == 0U) {
            ubBlockAligned = alignNum;// 至少保证一个对齐块
        }
        
        // 根据核数调整分块策略
        if (tiling->totalLength % blockDim == 0U) {
            // 均匀分块
            uint32_t blockLength = tiling->totalLength / blockDim;
            uint32_t tileNum, tileLength, lastTileLength;
            
            TilingParamsCalc(blockLength, ubBlockAligned, tileNum, tileLength, lastTileLength);
            
            tiling->blockLength = blockLength;
            tiling->tileNum = tileNum;
            tiling->tileLength = tileLength;
            tiling->lasttileLength = lastTileLength;
            tiling->isEvenCore = 1U;
        } else {
            // 非均匀分块
            uint32_t formerNum = tiling->totalLength % blockDim;
            uint32_t tailNum = blockDim - formerNum;
            
            uint32_t formerLength = (tiling->totalLength / blockDim) + 1;
            uint32_t tailLength = tiling->totalLength / blockDim;
            
            uint32_t formerTileNum, formerTileLength, formerLastTileLength;
            uint32_t tailTileNum, tailTileLength, tailLastTileLength;
            
            TilingParamsCalc(formerLength, ubBlockAligned, 
                           formerTileNum, formerTileLength, formerLastTileLength);
            TilingParamsCalc(tailLength, ubBlockAligned, 
                           tailTileNum, tailTileLength, tailLastTileLength);
            
            tiling->formerNum = formerNum;
            tiling->formerLength = formerLength;
            tiling->formerTileNum = formerTileNum;
            tiling->formerTileLength = formerTileLength;
            tiling->formerLastTileLength = formerLastTileLength;
            
            tiling->tailNum = tailNum;
            tiling->tailLength = tailLength;
            tiling->tailTileNum = tailTileNum;
            tiling->tailTileLength = tailTileLength;
            tiling->tailLastTileLength = tailLastTileLength;
            tiling->isEvenCore = 0U;
        }
    }
}